package org.liuyuantao.lucene.score;

import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.Date;

import org.apache.lucene.document.Document;
import org.apache.lucene.index.*;
import org.apache.lucene.queries.CustomScoreProvider;
import org.apache.lucene.queries.CustomScoreQuery;
import org.apache.lucene.queries.function.FunctionQuery;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.queries.function.valuesource.IntFieldSource;
import org.apache.lucene.search.*;
import org.apache.lucene.store.Directory;
import org.apache.lucene.util.BytesRef;

/**
 * 自定义评分
 */
public class CustomScoreQueryExample {

    private Directory directory;

    public CustomScoreQueryExample(Directory directory) {
        this.directory = directory;
    }

    /**
     * 自定义评分,是普通Query查询的评分除以FunctionQuery查询的评分的分数作为最后评分
     *
     * @param field
     * @param key
     */
    public void searchByScoreQuery(String field, String key) {
        try {
            IndexSearcher searcher = new IndexSearcher(DirectoryReader.open(directory));
            Query q = new TermQuery(new Term(field, key));
            //1、创建一个评分域
            ValueSource score = new IntFieldSource("score");
            FunctionQuery fd = new FunctionQuery(score);
            //2、根据评分域和原有的Query创建自定义的Query对象
            MyCustomScoreQuery query = new MyCustomScoreQuery(q, fd);
            TopDocs tds = null;
            tds = searcher.search(query, 100);
            SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd hh:mm:ss");
            for (ScoreDoc sd : tds.scoreDocs) {
                Document d = searcher.doc(sd.doc);
                System.out.println(sd.doc + ":(" + sd.score + ")" +
                        "[" + d.get("name") + "【" + d.get("path") + "】--->" +
                        d.get("size") + "-----" + sdf.format(new Date(Long.valueOf(d.get("date")))) + "]");
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public void searchByFileScoreQuery() {
        try {
            IndexSearcher searcher = new IndexSearcher(DirectoryReader.open(directory));
            Query q = new TermQuery(new Term("content", "java"));
            //1、创建一个评分域
            FilenameScoreQuery query = new FilenameScoreQuery(q);
            TopDocs tds = null;
            tds = searcher.search(query, 100);
            SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd hh:mm:ss");
            for (ScoreDoc sd : tds.scoreDocs) {
                Document d = searcher.doc(sd.doc);
                System.out.println(sd.doc + ":(" + sd.score + ")" +
                        "[" + d.get("name") + "【" + d.get("path") + "】--->" +
                        d.get("size") + "-----" + sdf.format(new Date(Long.valueOf(d.get("date")))) + "]");
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 根据文件名字评分排序
     */
    @SuppressWarnings("serial")
    private class FilenameScoreQuery extends CustomScoreQuery {

        public FilenameScoreQuery(Query subQuery) {
            super(subQuery);
        }

        @Override
        protected CustomScoreProvider getCustomScoreProvider(LeafReaderContext context) throws IOException {
            return new FilenameScoreProvider(context);
        }
    }

    /**
     * 根据日期进行评分排序
     */
    private class DateScoreProvider extends CustomScoreProvider {
        NumericDocValues dateList = null;

        public DateScoreProvider(LeafReaderContext context) throws IOException {
            super(context);
            LeafReader reader = context.reader();
            this.dateList = reader.getNumericDocValues("sortDate");
        }

        @Override
        public float customScore(int doc, float subQueryScore, float valSrcScore)
                throws IOException {
            long date = dateList.get(doc);
            long today = new Date().getTime();
            long year = 1000 * 60 * 60 * 24 * 365;
            if (today - date <= year) {
                //为其加分
            }
            return super.customScore(doc, subQueryScore, valSrcScore);
        }

    }

    private class FilenameScoreProvider extends CustomScoreProvider {
        BinaryDocValues nameList = null;

        /**
         * 此处nameList的值只能获取到DocValues类型的缓存(BinaryDocValues、NumericDocValues、SortedDocValues、SortedSetDocValues)
         *
         * @param context
         * @throws IOException
         */
        public FilenameScoreProvider(LeafReaderContext context) throws IOException {
            super(context);
            LeafReader reader = context.reader();
            nameList = reader.getBinaryDocValues("sortName");
        }

        @Override
        public float customScore(int doc, float subQueryScore, float valSrcScore)
                throws IOException {
            // 如何根据doc获取相应的field的值,在reader没有关闭之前,所有的数据会存储要一个域缓存中,可以通过域缓存获取很多有用的信息
            BytesRef bytesRef = nameList.get(doc);
            String filename = bytesRef.utf8ToString();
            if (filename.endsWith(".txt") || filename.endsWith(".ini")) {
                return subQueryScore * 1.5f;
            }
            return subQueryScore / 1.5f;
        }
    }

    /**
     * 默认情况实现的评分是通过原有的评分*传入进来的评分域所获取的评分来确定最终打分的
     * 为了根据不同的需求进行评分,需要自己进行评分的设定
     */
    @SuppressWarnings("serial")
    private class MyCustomScoreQuery extends CustomScoreQuery {
        public MyCustomScoreQuery(Query subQuery, FunctionQuery scoringQuery) {
            super(subQuery, scoringQuery);
        }

        /**
         * @param context
         * @return
         * @throws IOException
         */
        @Override
        protected CustomScoreProvider getCustomScoreProvider(LeafReaderContext context) throws IOException {

            /**
             * 自定评分的步骤
             * 创建一个类继承于CustomScoreProvider
             * 覆盖customScore方法
             */
            return new MyCustomScoreProvider(context);
        }

    }

    private class MyCustomScoreProvider extends CustomScoreProvider {

        public MyCustomScoreProvider(LeafReaderContext context) {
            super(context);
        }

        /**
         * 自定义评分规则
         *
         * @param doc
         * @param subQueryScore 默认文档的打分
         * @param valSrcScore   评分与的打分
         * @return
         * @throws IOException
         */
        @Override
        public float customScore(int doc, float subQueryScore, float valSrcScore) throws IOException {
            return subQueryScore / valSrcScore;
        }
    }

}
